import sys
from PyQt5.QtCore import Qt, QRegExp
from PyQt5.QtGui import QRegExpValidator
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, QFormLayout,
                             QPushButton, QLabel, QLineEdit, QFileDialog, QComboBox, QCheckBox, QTextEdit)
from ansys.fluent.core import launch_fluent


class ProcessTab(QWidget):
    def __init__(self, parent):
        super().__init__(parent)
        self.parent = parent
        self.mesh_file = None
        self.energy_model = "no"
        self.initUI()

    def initUI(self):
        layout = QVBoxLayout()

        form_layout = QFormLayout()
        form_layout.setSpacing(20)  # 设置表单项之间的间隔

        # 激活能量方程
        self.energyCheckBox = QCheckBox(self)
        self.energyCheckBox.stateChanged.connect(self.toggleEnergyModel)
        form_layout.addRow('激活能量方程:', self.energyCheckBox)

        # 设置湍流方程
        self.energyComboBox = QComboBox(self)
        equations = {
            'k-kl-w': 'k-kl-w',
            'ke_realizable': 'ke_realizable',
            'kw_bsl': 'kw_bsl',
            'kw_geko': 'kw_geko',
            'kw_sst': 'kw_sst',
            'kw_standard': 'kw_standard',
        }
        for key, value in equations.items():
            self.energyComboBox.addItem(value, userData=key)
        self.energyComboBox.currentIndexChanged.connect(self.viscous_q)
        form_layout.addRow('设置湍流方程:', self.energyComboBox)

        # 设置求解方法
        self.solveMethodComboBox = QComboBox(self)
        self.solveMethodComboBox.addItems(
            ['SIMPLE', 'PISO', 'SIMPLEC', 'Coupled'])
        self.solveMethodComboBox.currentIndexChanged.connect(self.methods)
        form_layout.addRow('设置求解方法:', self.solveMethodComboBox)

        # 设置求解步数
        self.solveStepsLineEdit = QLineEdit(self)
        self.solveStepsLineEdit.setValidator(
            QRegExpValidator(QRegExp("[0-9]*")))
        form_layout.addRow('设置求解步数:', self.solveStepsLineEdit)

        # 输入风速
        self.windSpeedLineEdit = QLineEdit(self)
        self.windSpeedLineEdit.setValidator(
            QRegExpValidator(QRegExp("[0-9]*")))
        form_layout.addRow('输入风速:', self.windSpeedLineEdit)

        layout.addLayout(form_layout)

        # 主函数处理按钮
        processButton = QPushButton("处理网格文件", self)
        processButton.setMinimumWidth(
            self.windSpeedLineEdit.minimumSizeHint().width())
        processButton.clicked.connect(self.process_mesh_file)

        # 添加按钮的居中布局
        button_layout = QHBoxLayout()
        button_layout.addStretch(1)
        button_layout.addWidget(processButton)
        button_layout.addStretch(1)

        layout.addLayout(button_layout)

        self.setLayout(layout)

        # 设置默认选项并触发相应函数
        self.energyComboBox.setCurrentText('k-kl-w')
        self.solveMethodComboBox.setCurrentText('SIMPLE')

    def setMeshFile(self, mesh_file):
        self.mesh_file = mesh_file

    def toggleEnergyModel(self, state):
        if state == Qt.Checked:
            self.energy_model = "yes"
            self.parent.consoleOutput.append(
                '<span style="color:green;">能量方程已激活！</span>')
        else:
            self.energy_model = "no"
            self.parent.consoleOutput.append(
                '<span style="color:red;">能量方程已关闭！</span>')

    def viscous_q(self):
        selected_equation = self.energyComboBox.currentText()
        if selected_equation == 'k-kl-w':
            self.parent.session.tui.define.models.viscous.k_kl_w('yes')
            self.parent.consoleOutput.append(
                '<span style="color:blue;">湍流方程 k-kl-w 已设置！</span>')
        elif selected_equation == 'ke_realizable':
            self.parent.session.tui.define.models.viscous.ke_realizable('yes')
            self.parent.consoleOutput.append(
                '<span style="color:blue;">湍流方程 ke_realizable 已设置！</span>')
        elif selected_equation == 'kw_bsl':
            self.parent.session.tui.define.models.viscous.kw_bsl('yes')
            self.parent.consoleOutput.append(
                '<span style="color:blue;">湍流方程 kw_bsl 已设置！</span>')
        elif selected_equation == 'kw_geko':
            self.parent.session.tui.define.models.viscous.kw_geko('yes')
            self.parent.consoleOutput.append(
                '<span style="color:blue;">湍流方程 kw_geko 已设置！</span>')
        elif selected_equation == 'kw_sst':
            self.parent.session.tui.define.models.viscous.kw_sst('yes')
            self.parent.consoleOutput.append(
                '<span style="color:blue;">湍流方程 kw_sst 已设置！</span>')
        elif selected_equation == 'kw_standard':
            self.parent.session.tui.define.models.viscous.kw_standard('yes')
            self.parent.consoleOutput.append(
                '<span style="color:blue;">湍流方程 kw_standard 已设置！</span>')

    def methods(self):
        selected_method = self.solveMethodComboBox.currentText()
        if selected_method == 'SIMPLE':
            self.parent.session.solution.methods.p_v_coupling.flow_scheme = 'SIMPLE'
            self.parent.consoleOutput.append(
                '<span style="color:purple;">求解方法 SIMPLE 已设置！</span>')
        elif selected_method == 'PISO':
            self.parent.session.solution.methods.p_v_coupling.flow_scheme = 'PISO'
            self.parent.consoleOutput.append(
                '<span style="color:purple;">求解方法 PISO 已设置！</span>')
        elif selected_method == 'SIMPLEC':
            self.parent.session.solution.methods.p_v_coupling.flow_scheme = 'SIMPLEC'
            self.parent.consoleOutput.append(
                '<span style="color:purple;">求解方法 SIMPLEC 已设置！</span>')
        elif selected_method == 'Coupled':
            self.parent.session.solution.methods.p_v_coupling.flow_scheme = 'Coupled'
            self.parent.consoleOutput.append(
                '<span style="color:purple;">求解方法 Coupled 已设置！</span>')

    def process_mesh_file(self):
        if self.mesh_file is None:
            self.parent.consoleOutput.append("没有网格文件可以处理")
            return

        wind_speed = self.windSpeedLineEdit.text()
        iteration_steps = self.solveStepsLineEdit.text()

        self.parent.consoleOutput.append(f"处理网格文件: {self.mesh_file}")
        self.parent.consoleOutput.append(f"设置风速为: {wind_speed}")
        self.parent.consoleOutput.append(f"设置迭代步数为: {iteration_steps}")

        # 启动 Fluent 求解器模式
        solver = launch_fluent(
            precision="double", processor_count=4, mode="solver", show_gui=True)

        # 读取网格文件
        solver.tui.file.read_case(case_file_name=self.mesh_file)

        # 设置能量模型
        solver.tui.define.models.energy(
            self.energy_model, ", ", ", ", ", ", ", ")

        # 定义边界条件，包括动态传参风速
        solver.tui.define.boundary_conditions.velocity_inlet(
            "inlet",
            "no",
            "no",
            "yes",
            "yes",
            "no",
            wind_speed,  # 动态传参风速
            "no",
            "0",
            "no",
            "no",
            "yes",
            "5",
            "10",
        )

        # 设置监控参数并初始化求解
        solver.tui.solve.monitors.residual.plot("yes")
        solver.tui.solve.initialize.hyb_initialization()

        # 执行求解迭代
        solver.tui.solve.iterate(int(iteration_steps))  # 动态传参迭代步数

        # 保存求解结果
        solver.tui.file.write_case("output.cas.h5")
        solver.tui.file.write_data("output.dat.h5")

        # 退出 Fluent
        solver.exit()

        self.parent.consoleOutput.append(
            '<span style="color:green;">求解完成，结果已保存！</span>')


# 用于测试独立运行
if __name__ == '__main__':
    app = QApplication(sys.argv)
    main = ProcessTab(None)
    main.show()
    sys.exit(app.exec_())
